import numpy as np
from sklearn.datasets import load_iris,load_digits
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelBinarizer

class BP:
    def __init__(self, input_num, output_num, layer_num,alpha = 0.1):
        """
        :param input_num: 输入层的神经元个数
        :param output_num: 输出层的神经元个数
        :param layer_num: 隐藏层的神经元个数
        """

        # 学习率
        self.alpha = alpha

        # 权重参数
        self.w1 = np.random.random((layer_num, input_num))*2-1
        self.w2 = np.random.random((output_num, layer_num))*2-1
        # 常数
        self.b1 = np.random.random((layer_num, 1))
        self.b2 = np.random.random((output_num, 1))

    def forward(self):
        """
        前推函数
        """
        self.z1 = self.w1 * self.x.T + self.b1
        self.o1 = self.Sigmoid(self.z1)
        self.z2 = self.w2 * self.o1 + self.b2
        self.o2 = self.Sigmoid(self.z2)

    def backward(self):
        """
        回退函数，更新参数
        """
        # 更新隐藏层与输出层之间的参数
        delta2 = self.Sigmoid_grad(self.z2) * (self.o2 - self.y.T)
        self.w2 -= self.alpha * (delta2 * self.o1.T)
        self.b2 -= self.alpha * delta2

        # 更新隐藏层与输入层之间的参数
        delta1 = self.Sigmoid_grad(self.z1) * (self.w2.T * delta2)
        self.w1 -= self.alpha * (delta1 * self.x)
        self.b1 -= self.alpha * delta1

    def train(self, x, y, iters = 500):
        """
        训练函数
        :param x: 训练的原始数据
        :param y: 每次训练输出的结果
        :param iters: 迭代次数
        """
        # 进行iter次训练
        for iter in range(iters):
            # 每次训练随机选一个数据进行训练，并且依此来调整参数
            i = np.random.randint(x.shape[0])
            self.x = x[i]
            self.y = y[i]
            self.forward()
            self.backward()

    def predict(self, x,y_test):
        """
        预测函数
        :param x: 测试集的原始数据
        :param y_test: 测试集的标准答案
        """
        # 存储预测结果
        ans = []
        # 对每个数据进行预测
        for i in range(x.shape[0]):
            self.x = x[i]
            # 直接前推，输出的y就是预测结果
            self.forward()
            # 将输出的数据进行处理，变成和标签一样的格式
            ans.append(np.argmax(self.o2))
        # 预测的正确率
        print(np.mean(np.equal(ans,y_test)))

    def Sigmoid(self, x):
        """
        激活函数
        :param x: 变量
        """
        return 1 / (1 + np.exp(-x))

    def Sigmoid_grad(self, x):
        """
        求导函数
        :param x: 变量
        """
        # 因为x是向量，所以求导结果为矩阵
        ans = np.eye(len(x))
        for i in range(len(x)):
            ans[i, i] = self.Sigmoid(x[i, 0]) * (1 - self.Sigmoid(x[i, 0]))
        return ans


if __name__ == '__main__':
    # 载入鸢尾花集
    iris = load_iris()
    # 分割数据集
    x_train, x_test, y_train, y_test = train_test_split(iris.data, iris.target)
    # 转换格式 方便计算
    x_train = np.mat(x_train)
    x_test = np.mat(x_test)
    y_train = np.mat(y_train).T

    # 标签二值化
    y_train = LabelBinarizer().fit_transform(y_train)
    y_train = np.mat(y_train)

    # 生成BP模型
    bp1 = BP(input_num=4, output_num=3, layer_num=6)
    # 进行训练
    bp1.train(x_train, y_train,20000)

    # 测试
    print("鸢尾花集预测结果：")
    bp1.predict(x_test,y_test)

    # 载入手写数字数据集
    digits = load_digits()
    X = digits.data
    Y = digits.target
    X -= X.min()
    X /= X.max()

    x_train, x_test, y_train, y_test = train_test_split(X, Y)
    # 转换格式 方便计算
    x_train = np.mat(x_train)
    x_test = np.mat(x_test)
    y_train = np.mat(y_train).T

    # 标签二值化
    y_train = LabelBinarizer().fit_transform(y_train)
    y_train = np.mat(y_train)

    # 生成BP模型
    bp2 = BP(input_num=x_train.shape[1], output_num=y_train.shape[1], layer_num=20,alpha=0.1)
    # 进行训练
    bp2.train(x_train, y_train,30000)

    # 测试
    print("手写数字数据集预测结果：")
    bp2.predict(x_test,y_test)